#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
IMU数据监控工具
用于分析IMU数据质量，检测导致y轴偏差的问题
"""

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Imu
from nav_msgs.msg import Odometry
import math
import numpy as np
import time
from collections import deque

class IMUMonitor(Node):
    def __init__(self):
        super().__init__('imu_monitor')
        
        # 订阅IMU数据
        self.imu_subscriber = self.create_subscription(
            Imu, '/imu/data', self.imu_callback, 10)
        
        # 订阅原始IMU数据（如果有的话）
        self.raw_imu_subscriber = self.create_subscription(
            Imu, '/imu', self.raw_imu_callback, 10)
            
        # 订阅里程计数据用于对比
        self.odom_subscriber = self.create_subscription(
            Odometry, '/odometry/filtered', self.odom_callback, 10)
        
        # 数据缓存
        self.imu_data = deque(maxlen=100)
        self.raw_imu_data = deque(maxlen=100)
        self.odom_data = deque(maxlen=100)
        
        # 统计信息
        self.imu_count = 0
        self.raw_imu_count = 0
        self.last_report_time = time.time()
        
        # 创建定时器定期输出报告
        self.timer = self.create_timer(5.0, self.print_analysis)
        
        self.get_logger().info("IMU监控工具启动")
    
    def imu_callback(self, msg):
        """处理滤波后的IMU数据"""
        self.imu_count += 1
        
        # 提取角速度和加速度
        angular_vel = msg.angular_velocity
        linear_acc = msg.linear_acceleration
        orientation = msg.orientation
        
        # 计算yaw角
        yaw = self.quaternion_to_yaw(orientation)
        
        data_point = {
            'timestamp': time.time(),
            'angular_vel_z': angular_vel.z,
            'linear_acc_x': linear_acc.x,
            'linear_acc_y': linear_acc.y,
            'linear_acc_z': linear_acc.z,
            'yaw': yaw
        }
        
        self.imu_data.append(data_point)
    
    def raw_imu_callback(self, msg):
        """处理原始IMU数据"""
        self.raw_imu_count += 1
        
        angular_vel = msg.angular_velocity
        linear_acc = msg.linear_acceleration
        
        data_point = {
            'timestamp': time.time(),
            'angular_vel_z': angular_vel.z,
            'linear_acc_x': linear_acc.x,
            'linear_acc_y': linear_acc.y,
            'linear_acc_z': linear_acc.z
        }
        
        self.raw_imu_data.append(data_point)
    
    def odom_callback(self, msg):
        """处理里程计数据"""
        twist = msg.twist.twist
        pose = msg.pose.pose
        
        yaw = self.quaternion_to_yaw(pose.orientation)
        
        data_point = {
            'timestamp': time.time(),
            'vx': twist.linear.x,
            'vy': twist.linear.y,
            'angular_vel_z': twist.angular.z,
            'x': pose.position.x,
            'y': pose.position.y,
            'yaw': yaw
        }
        
        self.odom_data.append(data_point)
    
    def quaternion_to_yaw(self, quaternion):
        """四元数转yaw角"""
        siny_cosp = 2 * (quaternion.w * quaternion.z + quaternion.x * quaternion.y)
        cosy_cosp = 1 - 2 * (quaternion.y * quaternion.y + quaternion.z * quaternion.z)
        return math.atan2(siny_cosp, cosy_cosp)
    
    def analyze_imu_bias(self, data_buffer):
        """分析IMU偏差"""
        if len(data_buffer) < 10:
            return None
            
        angular_vels = [d['angular_vel_z'] for d in data_buffer]
        acc_x = [d['linear_acc_x'] for d in data_buffer]
        acc_y = [d['linear_acc_y'] for d in data_buffer]
        acc_z = [d['linear_acc_z'] for d in data_buffer]
        
        # 静态时的偏差（假设最近的数据中有静态时刻）
        # 检测静态状态：角速度和x,y加速度都很小
        static_indices = []
        for i, d in enumerate(data_buffer):
            if (abs(d['angular_vel_z']) < 0.01 and 
                abs(d['linear_acc_x']) < 0.5 and 
                abs(d['linear_acc_y']) < 0.5):
                static_indices.append(i)
        
        if len(static_indices) > 5:
            static_angular_vels = [angular_vels[i] for i in static_indices]
            static_acc_x = [acc_x[i] for i in static_indices]
            static_acc_y = [acc_y[i] for i in static_indices]
            
            bias_analysis = {
                'angular_vel_bias': np.mean(static_angular_vels),
                'angular_vel_noise': np.std(static_angular_vels),
                'acc_x_bias': np.mean(static_acc_x),
                'acc_y_bias': np.mean(static_acc_y),
                'acc_x_noise': np.std(static_acc_x),
                'acc_y_noise': np.std(static_acc_y),
                'static_samples': len(static_indices)
            }
        else:
            bias_analysis = {
                'angular_vel_bias': np.mean(angular_vels),
                'angular_vel_noise': np.std(angular_vels),
                'acc_x_bias': np.mean(acc_x),
                'acc_y_bias': np.mean(acc_y),
                'acc_x_noise': np.std(acc_x),
                'acc_y_noise': np.std(acc_y),
                'static_samples': 0
            }
        
        return bias_analysis
    
    def print_analysis(self):
        """打印分析报告"""
        current_time = time.time()
        time_diff = current_time - self.last_report_time
        
        print("\n" + "="*60)
        print(f"IMU监控报告 - {time.strftime('%H:%M:%S')}")
        print("="*60)
        
        # 数据接收统计
        print(f"过去{time_diff:.1f}秒数据接收:")
        print(f"  滤波IMU数据: {self.imu_count}条")
        print(f"  原始IMU数据: {self.raw_imu_count}条")
        print(f"  里程计数据: {len(self.odom_data)}条")
        
        # 分析滤波后的IMU数据
        if len(self.imu_data) > 10:
            print("\n滤波后IMU数据分析:")
            bias_analysis = self.analyze_imu_bias(self.imu_data)
            if bias_analysis:
                print(f"  角速度偏差: {bias_analysis['angular_vel_bias']:.6f} rad/s")
                print(f"  角速度噪声: {bias_analysis['angular_vel_noise']:.6f} rad/s")
                print(f"  X加速度偏差: {bias_analysis['acc_x_bias']:.3f} m/s²")
                print(f"  Y加速度偏差: {bias_analysis['acc_y_bias']:.3f} m/s²")
                print(f"  X加速度噪声: {bias_analysis['acc_x_noise']:.3f} m/s²")
                print(f"  Y加速度噪声: {bias_analysis['acc_y_noise']:.3f} m/s²")
                print(f"  静态样本数: {bias_analysis['static_samples']}")
                
                # 提供建议
                print("\n问题分析和建议:")
                if abs(bias_analysis['angular_vel_bias']) > 0.01:
                    print(f"  ⚠️  角速度偏差过大 ({bias_analysis['angular_vel_bias']:.6f})")
                    print("     建议: 重新校准IMU或在EKF中增加角速度拒绝阈值")
                
                if bias_analysis['angular_vel_noise'] > 0.02:
                    print(f"  ⚠️  角速度噪声过大 ({bias_analysis['angular_vel_noise']:.6f})")
                    print("     建议: 降低IMU在EKF中的权重")
                
                if abs(bias_analysis['acc_y_bias']) > 0.5:
                    print(f"  ⚠️  Y轴加速度偏差过大 ({bias_analysis['acc_y_bias']:.3f})")
                    print("     建议: 检查IMU安装方向或禁用加速度融合")
        
        # 对比IMU和里程计的角速度
        if len(self.imu_data) > 5 and len(self.odom_data) > 5:
            print("\nIMU vs 里程计角速度对比:")
            recent_imu = list(self.imu_data)[-5:]
            recent_odom = list(self.odom_data)[-5:]
            
            imu_angular_vel = np.mean([d['angular_vel_z'] for d in recent_imu])
            odom_angular_vel = np.mean([d['angular_vel_z'] for d in recent_odom])
            
            print(f"  IMU角速度(平均): {imu_angular_vel:.6f} rad/s")
            print(f"  里程计角速度(平均): {odom_angular_vel:.6f} rad/s")
            print(f"  差值: {imu_angular_vel - odom_angular_vel:.6f} rad/s")
            
            if abs(imu_angular_vel - odom_angular_vel) > 0.05:
                print("  ⚠️  IMU和里程计角速度差异过大")
                print("     可能导致EKF融合时产生y轴偏差")
        
        self.last_report_time = current_time
        self.imu_count = 0
        self.raw_imu_count = 0

def main():
    rclpy.init()
    
    monitor = IMUMonitor()
    
    try:
        print("IMU监控工具启动，按Ctrl+C停止...")
        rclpy.spin(monitor)
    except KeyboardInterrupt:
        print("\n监控停止")
    finally:
        monitor.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main() 